from .Utils import get_rotation_mat, get_rotation_mat_interp

from vtkmodules.vtkCommonCore import vtkAnimationCue

from vtkmodules.vtkCommonDataModel import (
    vtkAnimationScene,
)

import numpy as np
import math

# acceleration by numba
from numba import jit

@jit(nopython=True)
def _get_interp_idx(t, x, idx):
    idx = np.interp(t, x, idx)
    idx0 = int(math.floor(idx))
    idx1 = int(math.ceil(idx))
    f = idx - idx0
    return idx0, idx1, f

# define callback for animation cue
class AnimationChannel(vtkAnimationCue):
    def __init__(self, node, path, input, output, interpolation) -> None:
        
        self.node = node
        self.path = path
        self.input = input
        self.output = output
        self.interpolation = interpolation
        self.index = np.arange(0, len(self.input-1), 1)

        self.SetStartTime(input[0])
        self.SetEndTime(input[-1])
        
        self.AddObserver("StartAnimationCueEvent",   self.start)
        self.AddObserver("EndAnimationCueEvent",     self.end)
        self.AddObserver("AnimationCueTickEvent",    self.tick)

    def apply(self, t):    
        idx0, idx1, f = _get_interp_idx(t, self.input, self.index)
        if self.path == 'rotation':
            r = get_rotation_mat_interp(self.output[idx0], self.output[idx1], f)
            self.node.rotate.SetMatrix(r)
        elif self.path == 'translation' or self.path == 'scale':
            v = self.output[idx0] * (1-f) + self.output[idx1] * f          
            if self.path == 'translation':                  
                self.node.translate.Identity()
                self.node.translate.Translate(v)
            elif self.path == 'scale': 
                self.node.scale.Identity()
                self.node.scale.Scale(v)
        else:
            print("Unimplemented animation path:", self.path)


    def setIdx(self, idx):
        if self.path == 'translation': 
            self.node.translate.Identity()            
            self.node.translate.Translate(self.output[idx]) 
        elif self.path == 'rotation':
            self.node.rotate.SetMatrix(get_rotation_mat(self.output[idx]))  
        elif self.path == 'scale':
            self.node.scale.Identity()
            self.node.scale.Scale(self.output[idx])


    def start(self, info, event):
        self.setIdx(0)

    def end(self, info, event):
        self.setIdx(-1)              

    def tick(self, info, event):
        self.apply(info.GetAnimationTime()) 


class Animation(vtkAnimationScene):
    def __init__(self, name:str="Animation") -> None:
        super().__init__()
        self.name = name
        self.play = False